#ifndef RAINFALLRUNOFF_HPP
#define RAINFALLRUNOFF_HPP

#include "edgeTraversal.hpp"
#include "sweSolver.hpp"
#include "solverBase.hpp"
#include "meshTypes.hpp"
#include "meshFunctions.hpp"
#include "magic_enum.hpp"
#include "vtkFunctions_export.hpp"

#include "solverSettings.hpp"

#include <filesystem>
namespace fs = std::filesystem;


using namespace hms;

struct rainfallRunoff{

	hms::Time time;
	hms::Traversal traversal;
	hms::EdgeTraversal edgeTraversal;
	hms::CellTraversal cellTraversal;
	Index blocksize;
	std::string fluxSchemeStr;
	bool
		useMuscl,
		useExport,
		useAtomicWrite;
	scalar maxCourant;
	fs::path exportPath;

	rainfallRunoff( hms::Time&& t, hms::SolverSettings&& s ) :
		time { std::move(t) },
		traversal { s.traversal },
		edgeTraversal { s.edgeTraversal },
		cellTraversal { s.cellTraversal },
		blocksize { s.blocksize },
		fluxSchemeStr { std::move(s.fluxScheme) },
		// useFriction { s.useFriction },
		// useSourceFields { s.useSourceFields },
		useMuscl { s.useMuscl },
		useExport { s.useExport },
		useAtomicWrite { s.useAtomicWrite },
		maxCourant { s.maxCourant },
		exportPath { std::move(s.exportPath) }
	{}
	
	template<typename MeshType>
	void operator()(MeshType& mesh){
		
		/**************************
		***** case properties *****
		***************************/

		bool
			useFriction          {true},
			useSourceFields      {true};
		scalar 
			manningCatchment     {0.015},
			manningChannel       {0.15},
			channelLength        {1000.0},
			channelWidth         {20.0},
			channelSlope         {0.02},
			catchmentSlope       {0.05},
			channelMinDepth      {1.0},
			rainfallSource       {3e-6}, // <-- [m/s]
			sourceEndTime        {5400.0};
		LimiterFunction lf       {LimiterFunction::minMod};

		/**************************
		*********** mesh **********
		***************************/
		mesh.blocksize(blocksize);
		Index
			nx { mesh.nCellsX() },
			ny { mesh.nCellsY() },
			ng { mesh.nGhostLayers() };


		static_assert( std::is_same_v<MeshType, UniMesh> || std:: is_same_v<MeshType, RectMesh>, 
		"Rainfall Runoff Verification case can only run on MeshTypes: RectMesh, UniMesh. The entered type does not match. Please choose on of the two options.");
		
		/**************************
		*********** Time **********
		***************************/
		// Time time;

		// time.startAt(0.0);
		// time.endAt(simulationTime);
		// time.maxTimeStep(maxTimeStep);
		// time.writeAtTimeIntervals(useWriteTimeInterval);
		// if ( useWriteTimeInterval ) time.writeTimeInterval(writeInterval);
		// else time.writeStepInterval(static_cast<hms::Index>(writeInterval));

		/**************************
		**** solver properties ****
		***************************/
		SWESolver<MeshType> solver {mesh, time};

		using FluxScheme = typename SolverTraits<SWESolver<MeshType>>::FluxScheme;
		FluxScheme fluxScheme = FluxScheme::HLLC;

		if(!fluxSchemeStr.empty()) {
			auto fluxSchemeOpt { magic_enum::enum_cast<FluxScheme>(fluxSchemeStr) };
			if ( fluxSchemeOpt.has_value() )
			fluxScheme = fluxSchemeOpt.value();
		}
		else {
			std::cerr << "WARNING: Flux Scheme \"" << fluxSchemeStr
				<< "\" unknown! Defaulting to HLLC\n";
		fluxScheme = FluxScheme::HLLC;
		}

		solver.useAtomicWrite(useAtomicWrite);
		solver.fluxScheme(fluxScheme);
		solver.traversal(traversal);
		solver.edgeTraversal(edgeTraversal);
		solver.cellTraversal(cellTraversal);
		solver.useFriction(useFriction);
		solver.useSourceFields(useSourceFields);
		solver.useMuscl(useMuscl);
		solver.limiterFunction(lf);
		solver.maxCourant(maxCourant);

		/**************************
		******** fields ***********
		***************************/
		solver.v().values() = 0.0;
		solver.d().values() = 0.0;
		solver.z().values() = 1.0;

		/**************************
		******** friction *********
		***************************/
		// construct an array which shares the same size as the mesh. 
		ArrayXXs frictionValues ((ny+2*ng), (nx+2*ng));
		
		// assign catchment roughness to complete array.
		frictionValues = manningCatchment;
		
		// assign channel roughness to the two middle rows.
		frictionValues.middleRows(frictionValues.rows()/2 - 1, 2) = manningChannel;
		
		// reshaped and transpose the array to a 1D row array and assign it to the friction field.
		solver.n() = ConstantField<NonUniform, MeshType>{ 
			mesh, 
			std::move(frictionValues.reshaped<Eigen::RowMajor>().transpose()) 
		};

		// same procedure as in friction assignment, but only in domain cells!
		ArrayXXs rainfallValues (ny, nx);
		rainfallValues = rainfallSource;
		rainfallValues.middleRows(rainfallValues.rows()/2 -1, 2) = 0.0;
		solver.r() =SourceField<NonUniform, MeshType>{ 
			mesh,
			{
				{0, rainfallValues.reshaped<Eigen::RowMajor>().transpose()},
				{sourceEndTime, ArrayXXs::Zero(1, mesh.nDomainCells())}
			},
			time
		};
		// no infiltration
		solver.i() = SourceField<Uniform, MeshType>{ mesh, ArrayNs<1>{0} };

		/**************************
		**** boundary patches *****
		***************************/
		// erase the right patch and create three new patches for the right side
		// two patches along the sloped catchment and one patch for the channel outflow.
		mesh.eraseBoundaryPatch("right");

		// lower catchment side boundary patch
		Index 
			size  {ny / 2 - 1},
			begin {nx * (ny + 1) + nx },
			end   {nx * (ny + 1) 
					+ (nx+1) * (size-1) 
					+ nx };

		mesh.addBoundaryPatch( "rightLowerCatchment",
			ArrayXi::LinSpaced( size, begin, end ) );

		// channel boundary patch
		size  = 2;
		begin = nx * (ny + 1) 
			+ (nx+1) * (ny / 2 - 1) 
			+ nx;
		end   = nx * (ny + 1) 
			+ (nx+1) * (ny / 2) 
			+ nx;

		mesh.addBoundaryPatch( "rightChannel", 
			ArrayXi::LinSpaced(size, begin, end) );

		// upper catchment side boundary patch
		size  = ny / 2 - 1;
		begin = nx * (ny + 1) 
			+ (nx+1) * (size+2) 
			+ nx;
		end = nx * (ny + 1)
			+ (nx+1) * (ny-1)
			+ nx;

		mesh.addBoundaryPatch( "rightUpperCatchment", 
			ArrayXi::LinSpaced(size, begin, end) );

		// get patches, except right patch
		const BoundaryPatch& left   { mesh.boundaryPatch("left") };
		const BoundaryPatch& bottom { mesh.boundaryPatch("bottom") };
		const BoundaryPatch& top    { mesh.boundaryPatch("top") };
		const BoundaryPatch& rightLowerCatchment { mesh.boundaryPatch("rightLowerCatchment") };
		const BoundaryPatch& rightChannel        { mesh.boundaryPatch("rightChannel") };
		const BoundaryPatch& rightUpperCatchment { mesh.boundaryPatch("rightUpperCatchment") };

		/**************************
		**** bottom elevation *****
		***************************/
		mesh.forEachCell( [&](const auto& cell){
		
			scalar x {cell.centroid()[0] - 500.0};
			scalar y {cell.centroid()[1] - 810.0};
			scalar z {0.0};

			// on the one hillside
			if(y > 0.5 * channelWidth) {
				z += y * catchmentSlope;
				if(x > 0.5 * channelLength) z += 0.1;
			}

			// on the other hillside
			if(y < -0.5 * channelWidth) {
				z -= y * catchmentSlope;
				if(x > 0.5 * channelLength) z += 0.1;
			}

			// in the channel
			if(y < 0.5 * channelWidth && y > -0.5* channelWidth) 
			z = 0.5 * channelWidth * catchmentSlope 
				- channelMinDepth 
				- x * channelSlope 
				- 0.5 * channelLength * channelSlope;

			solver.z().values()(0, cell.index()) = z;
		});

		/**************************
		*** boundary conditions ***
		***************************/

		// get bottom elevation from boundary domain cells in order to use them as static values for boundary conditions.
		ArrayNs<1> bottomElevationTop    {solver.z().domainValues(top)(0,0)};
		ArrayNs<1> bottomElevationBottom {solver.z().domainValues(bottom)(0,0)};
		Array1Xs   bottomElevationLeft   {solver.z().domainValues(left)};
		Array1Xs   bottomElevationRUC    {solver.z().domainValues(rightUpperCatchment)};
		Array1Xs   bottomElevationRLC    {solver.z().domainValues(rightLowerCatchment)};

		// bed elevation in ghost cells at the channel will continue the slope.
		ArrayNs<1> ghostBedElevChannel;
		ghostBedElevChannel = -20.6;

		solver.z().setBoundaryCondition( left,   StaticValue{}, bottomElevationLeft,   false );
		solver.z().setBoundaryCondition( top,    StaticValue{}, bottomElevationTop,    true );
		solver.z().setBoundaryCondition( bottom, StaticValue{}, bottomElevationBottom, true );
		solver.z().setBoundaryCondition( rightLowerCatchment, StaticValue{}, bottomElevationRLC, false );
		solver.z().setBoundaryCondition( rightUpperCatchment, StaticValue{}, bottomElevationRUC, false );
		// ghost cells at the outflow continue the slope. last domain cells are -20.4
		solver.z().setBoundaryCondition( rightChannel,        StaticValue{}, ghostBedElevChannel, true ); 
		// water depth
		solver.d().setBoundaryCondition( left,   ZeroGradient{} );
		solver.d().setBoundaryCondition( top,    ZeroGradient{} );
		solver.d().setBoundaryCondition( bottom, ZeroGradient{} );
		solver.d().setBoundaryCondition( rightLowerCatchment, ZeroGradient{} );
		solver.d().setBoundaryCondition( rightUpperCatchment, ZeroGradient{} );
		solver.d().setBoundaryCondition( rightChannel,        FreeOutflow_waterdepth{} );
		// flow velocity
		solver.v().setBoundaryCondition( left,   Slip_velocity{} );
		solver.v().setBoundaryCondition( top,    Slip_velocity{} );
		solver.v().setBoundaryCondition( bottom, Slip_velocity{} );
		solver.v().setBoundaryCondition( rightLowerCatchment, Slip_velocity{} );
		solver.v().setBoundaryCondition( rightUpperCatchment, Slip_velocity{} );
		solver.v().setBoundaryCondition( rightChannel,        FreeOutflow_velocity{} );
		
		/**************************
		***** initialisation ******
		***************************/
		if(traversal == hms::Traversal::edges) {
			std::cout << "Edge traversal: "
				<< magic_enum::enum_name( solver.edgeTraversal() );
		} else {
			std::cout << "Cell traversal: "
				<< magic_enum::enum_name( solver.cellTraversal() );
		}

		std::cout << "Flux scheme: " << magic_enum::enum_name( solver.fluxScheme() ) <<std::endl;

		solver.initialise(); 
		mesh.printInfo();

		Timer t;
		std::cout << std::setprecision(5) << std::fixed;
		std::cout << "Time: " << time.current() << "s" << std::endl;

		std::unique_ptr<vtk::xml::XMLDocument> bindingFile {(useExport) ? vtk::createBindingDoc() : nullptr };
		// rain stops after sourceEndTime is reached.
		bool rainHasStopped {false};
		while ( !time.atEnd()){
			// stop the rain, when the time is right. maybe this does not work. if so, try values()-function of r(). 
			if( time.current() > sourceEndTime && !rainHasStopped){
				// solver.r() = SourceField<Uniform, MeshType>{ mesh, ArrayNs<1>{0} };
				solver.useSourceFields(false);
				rainHasStopped = true;
			}
			// vtk export
			if( useExport && time.atWriteTime() ) {
				vtk::toVtk(exportPath,mesh,solver,time.nWritten());
				vtk::updateBinding(bindingFile.get(), mesh, time.current(), time.nWritten(),exportPath);
			}
			// evolution in time.
			solver.advance( time );
		}
		// last vtk export
		if( useExport ) {
			vtk::toVtk( exportPath,mesh,solver,time.nWritten() );
			vtk::updateBinding(bindingFile.get(), mesh, time.current(), time.nWritten(),exportPath);
		}

		double elapsed { t.elapsed() };

		// Timing
		double execTimeSum {0};
		for ( auto execTime : solver.execTimings() )
			execTimeSum += execTime;

		std::cout
			<< "\nComputation time: " << elapsed << "s"
			<< "\nSum of execution times logged by solver: " << execTimeSum << "s"
			<< "\n\t" << 100 * solver.execTimings()[0] / execTimeSum
				<< "% for time step calculation"
				<< " (" << solver.execTimings()[0] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[1] / execTimeSum
				<< "% for fluxes (including slope)"
				<< " (" << solver.execTimings()[1] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[2] / execTimeSum
				<< "% for sources " << ( useFriction ?
					std::string{ "(including friction)" } :
					std::string{ "(friction disabled)" } )
				<< " (" << solver.execTimings()[2] << "s)"	
			<< "\n\t" << 100 * solver.execTimings()[3] / execTimeSum
				<< "% for new state"
				<< " (" << solver.execTimings()[3] << "s)"
			<< "\n\t" << 100 * solver.execTimings()[4] / execTimeSum
				<< "% for updating boundaries"
				<< " (" << solver.execTimings()[4] << "s)"
			<< "\n";
		}
};


#endif